from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from typing import Tuple
import os, getpass, argparse, time
import yaml

class Surfer:
    def __init__(self, config_file:str="config.yml"):
        with open(config_file, "r") as f:
            default_user_data_dir = "/home/{}/.config/google-chrome".format(getpass.getuser())
            cfg = yaml.safe_load(f.read())

            # When an item is missing, use default value.
            self.config = {
                "USER_DATA_DIR"               : cfg.get("USER_DATA_DIR", default_user_data_dir),
                "PAGE_LOAD_TIMEOUT"           : cfg.get("PAGE_LOAD_TIMEOUT", 5),
                "CHROMEDRIVER_EXECUTABLE_PATH": cfg.get("CHROMEDRIVER_EXECUTABLE_PATH","chromedriver"),
                "HEADLESS"                    : cfg.get("HEADLESS", False)
            }
        os.environ['WDM_LOG_LEVEL'] = '0' # disable webdriver logging
        options = webdriver.ChromeOptions()
        options.add_argument("--user-data-dir=" + self.config["USER_DATA_DIR"])
        options.add_argument("--no-sandbox") # root user uncomment

        if self.config["HEADLESS"]:
            options.add_argument("--headless")
            options.add_argument("--disable-gpu")

        self.driver = webdriver.Chrome(
            service = Service(self.config["CHROMEDRIVER_EXECUTABLE_PATH"]),
            options = options
        )

        self.driver.set_page_load_timeout(self.config["PAGE_LOAD_TIMEOUT"])
       
    def visit(self, url_or_domain: str) -> Tuple[int, str]:
        """Visit a URL.

        Args:
            url_or_domain: URL or domain.

        Returns:
            status: Status code, 0: OK, other means error.
            info: Failure related infomation.
        """
        status, info = 0, ""

        # cover domain to URL
        url = url_or_domain
        if len(url_or_domain) <= 8 or url_or_domain[:8] != "https://":
            url = "https://" + url_or_domain

        try:
            self.driver.get(url)
            status = 0
        except Exception as e:
            status = 1
            info = str(e)

        return status, info

    def quit(self):
        self.driver.quit()

def main():
    parser = argparse.ArgumentParser(description="Manipulate the browser to browse the web.")

    parser.add_argument(
        "--config", "-c",
        type = str,
        default = "config.yml",
        help = "Config file path. (Default 'config.yml')")

    parser.add_argument(
        "--domain", "-d",
        type = str,
        default = "baidu.com",
        help = "Domain to visit.")

    args = parser.parse_args()

    surfer = Surfer(args.config)
    surfer.visit(args.domain)
    time.sleep(4)
    surfer.quit()

if __name__ == '__main__':
    main()